//! 线程池管理模块

use std::sync::{mpsc, Arc, Mutex};
use std::thread;
use crate::error::{DataForgeError, Result};

type Job = Box<dyn FnOnce() + Send + 'static>;

/// 自定义线程池
pub struct ThreadPool {
    workers: Vec<Worker>,
    sender: Option<mpsc::Sender<Job>>,
}

impl ThreadPool {
    /// 创建新的线程池
    pub fn new(size: usize) -> Result<ThreadPool> {
        if size == 0 {
            return Err(DataForgeError::validation("Thread pool size must be greater than 0"));
        }

        let (sender, receiver) = mpsc::channel();
        let receiver = Arc::new(Mutex::new(receiver));
        let mut workers = Vec::with_capacity(size);

        for id in 0..size {
            workers.push(Worker::new(id, Arc::clone(&receiver))?);
        }

        Ok(ThreadPool {
            workers,
            sender: Some(sender),
        })
    }

    /// 执行任务
    pub fn execute<F>(&self, f: F) -> Result<()>
    where
        F: FnOnce() + Send + 'static,
    {
        let job = Box::new(f);
        
        if let Some(sender) = &self.sender {
            sender.send(job)
                .map_err(|_| DataForgeError::generator("Failed to send job to thread pool"))?;
        } else {
            return Err(DataForgeError::generator("Thread pool has been shut down"));
        }

        Ok(())
    }

    /// 获取工作线程数量
    pub fn size(&self) -> usize {
        self.workers.len()
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        drop(self.sender.take());

        for worker in &mut self.workers {
            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap();
            }
        }
    }
}

/// 工作线程
struct Worker {
    #[allow(dead_code)]
    id: usize,
    thread: Option<thread::JoinHandle<()>>,
}

impl Worker {
    fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Job>>>) -> Result<Worker> {
        let thread = thread::Builder::new()
            .name(format!("dataforge-worker-{}", id))
            .spawn(move || loop {
                let message = receiver.lock().unwrap().recv();

                match message {
                    Ok(job) => {
                        job();
                    }
                    Err(_) => {
                        break;
                    }
                }
            })
            .map_err(|e| DataForgeError::generator(&format!("Failed to spawn worker thread: {}", e)))?;

        Ok(Worker {
            id,
            thread: Some(thread),
        })
    }
}

/// 线程池统计信息
#[derive(Debug, Clone)]
pub struct ThreadPoolStats {
    /// 活跃线程数
    pub active_threads: usize,
    /// 总线程数
    pub total_threads: usize,
    /// 队列中的任务数
    pub queued_jobs: usize,
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::time::Duration;

    #[test]
    fn test_thread_pool_creation() {
        let pool = ThreadPool::new(4);
        assert!(pool.is_ok());
        assert_eq!(pool.unwrap().size(), 4);
    }

    #[test]
    fn test_thread_pool_zero_size() {
        let pool = ThreadPool::new(0);
        assert!(pool.is_err());
    }

    #[test]
    fn test_thread_pool_execution() {
        let pool = ThreadPool::new(2).unwrap();
        let counter = Arc::new(AtomicUsize::new(0));
        
        for _ in 0..10 {
            let counter = Arc::clone(&counter);
            pool.execute(move || {
                counter.fetch_add(1, Ordering::SeqCst);
            }).unwrap();
        }

        // 等待任务完成
        thread::sleep(Duration::from_millis(100));
        assert_eq!(counter.load(Ordering::SeqCst), 10);
    }
}